from torchvision.models import resnet34
from torch import nn


class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.upch = nn.Sequential(
            nn.Conv2d(kernel_size=1, in_channels=1, out_channels=3, bias=True, stride=1),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True)
        )
        self.backbone = resnet34(pretrained=True)
        self.backbone.fc = nn.Linear(512, 2, bias=True)

    def forward(self, x):
        x = self.upch(x)
        x = self.backbone(x)
        return x


if __name__ == '__main__':
    import torch
    model = model()
    out = model(torch.randn(5, 1, 512, 512))
    print(out.size())